package hbbft

import (
	"bytes"
	"encoding/gob"
	"fmt"
	"sync"
)

// HoneyBadger Block处理
type HoneyBadgerBlock struct {
	Config
	// 阈值加密公钥
	TPKEPK *TPKEPublicKey
	// 阈值加密私钥
	TPKESK *TPKEPrivateKey
	// 接收到的解密份额
	decryptionShares map[uint64]map[int]*DecryptionShare // round -> sender -> share
	// 输出队列
	outputQueues map[uint64]chan [][]byte
	// 锁
	lock sync.RWMutex
	// 消息通道
	messageCh chan blockMessage
	// 关闭通道
	closeCh chan struct{}
}

type blockMessage struct {
	senderID uint64
	round    uint64
	shares   []*DecryptionShare
}

// 加密的事务数据
type EncryptedTransaction struct {
	TKey       *Ciphertext // 加密的密钥
	Ciphertext []byte      // 加密的事务数据
}

// 创建HoneyBadger Block处理器
func NewHoneyBadgerBlock(cfg Config, tpkepk *TPKEPublicKey, tpkesk *TPKEPrivateKey) *HoneyBadgerBlock {
	hbb := &HoneyBadgerBlock{
		Config:           cfg,
		TPKEPK:           tpkepk,
		TPKESK:           tpkesk,
		decryptionShares: make(map[uint64]map[int]*DecryptionShare),
		outputQueues:     make(map[uint64]chan [][]byte),
		messageCh:        make(chan blockMessage, 1024),
		closeCh:          make(chan struct{}),
	}
	
	go hbb.run()
	return hbb
}

// 运行HoneyBadger Block处理器
func (hbb *HoneyBadgerBlock) run() {
	for {
		select {
		case msg := <-hbb.messageCh:
			hbb.handleMessage(msg.senderID, msg.round, msg.shares)
		case <-hbb.closeCh:
			return
		}
	}
}

// 处理消息
func (hbb *HoneyBadgerBlock) handleMessage(senderID, round uint64, shares []*DecryptionShare) {
	hbb.lock.Lock()
	defer hbb.lock.Unlock()
	
	// 初始化轮次
	if hbb.decryptionShares[round] == nil {
		hbb.decryptionShares[round] = make(map[int]*DecryptionShare)
	}
	
	// 存储解密份额
	for _, share := range shares {
		if share != nil {
			hbb.decryptionShares[round][int(senderID)] = share
		}
	}
	
	// 检查是否达到阈值
	if len(hbb.decryptionShares[round]) >= hbb.F+1 {
		// 尝试解密所有加密事务
		hbb.tryDecrypt(round)
	}
}

// 尝试解密
func (hbb *HoneyBadgerBlock) tryDecrypt(round uint64) {
	// 这里需要与ACS的输出配合
	// 暂时返回空结果
	if hbb.outputQueues[round] != nil {
		select {
		case hbb.outputQueues[round] <- [][]byte{}:
		default:
		}
	}
}

// 处理加密事务
func (hbb *HoneyBadgerBlock) ProcessEncryptedTransactions(round uint64, encryptedTxs []*EncryptedTransaction) {
	// 为每个加密事务生成解密份额
	shares := make([]*DecryptionShare, len(encryptedTxs))
	
	for i, encTx := range encryptedTxs {
		if encTx != nil && encTx.TKey != nil {
			// 生成解密份额
			share := hbb.TPKESK.DecryptShare(encTx.TKey)
			shares[i] = share
		}
		_ = i // 避免未使用变量警告
	}
	
	// 广播解密份额
	hbb.BroadcastShares(round, shares)
}

// 广播解密份额
func (hbb *HoneyBadgerBlock) BroadcastShares(round uint64, shares []*DecryptionShare) *DecryptionSharesMessage {
	return &DecryptionSharesMessage{
		Round: round,
		Shares: shares,
	}
}

// 处理接收到的解密份额
func (hbb *HoneyBadgerBlock) HandleShares(senderID uint64, msg *DecryptionSharesMessage) {
	hbb.messageCh <- blockMessage{
		senderID: senderID,
		round:    msg.Round,
		shares:   msg.Shares,
	}
}

// 等待解密结果
func (hbb *HoneyBadgerBlock) WaitForDecryption(round uint64) [][]byte {
	hbb.lock.Lock()
	
	// 创建输出队列
	if hbb.outputQueues[round] == nil {
		hbb.outputQueues[round] = make(chan [][]byte, 1)
	}
	
	hbb.lock.Unlock()
	
	// 等待结果
	select {
	case result := <-hbb.outputQueues[round]:
		return result
	}
}

// 解密份额消息
type DecryptionSharesMessage struct {
	Round  uint64
	Shares []*DecryptionShare
}

// 关闭HoneyBadger Block处理器
func (hbb *HoneyBadgerBlock) Close() {
	close(hbb.closeCh)
}

// 加密事务
func EncryptTransactions(transactions []Transaction, tpkepk *TPKEPublicKey) ([]*EncryptedTransaction, error) {
	encryptedTxs := make([]*EncryptedTransaction, len(transactions))
	
	for i, tx := range transactions {
		// 序列化事务
		var buf bytes.Buffer
		encoder := gob.NewEncoder(&buf)
		if err := encoder.Encode(tx); err != nil {
			return nil, fmt.Errorf("failed to encode transaction: %v", err)
		}
		
		txData := buf.Bytes()
		
		// 生成随机密钥
		key := randBytes(32)
		if len(key) != 32 {
			return nil, fmt.Errorf("generated key has wrong length: %d", len(key))
		}
		
		// 使用AES加密事务数据
		ciphertext, err := AESEncrypt(key, txData)
		if err != nil {
			return nil, fmt.Errorf("failed to encrypt transaction: %v", err)
		}
		
		// 使用阈值加密加密密钥
		tkey, err := tpkepk.Encrypt(key)
		if err != nil {
			return nil, fmt.Errorf("failed to encrypt key: %v", err)
		}
		
		encryptedTxs[i] = &EncryptedTransaction{
			TKey:       tkey,
			Ciphertext: ciphertext,
		}
	}
	
	return encryptedTxs, nil
}

// 解密事务
func DecryptTransactions(encryptedTxs []*EncryptedTransaction, tpkepk *TPKEPublicKey, shares map[int]*DecryptionShare) ([][]byte, error) {
	decryptedTxs := make([][]byte, len(encryptedTxs))
	
	for i, encTx := range encryptedTxs {
		if encTx == nil {
			continue
		}
		
		// 组合解密份额恢复密钥
		key, err := tpkepk.CombineShares(encTx.TKey, shares)
		if err != nil {
			return nil, fmt.Errorf("failed to combine shares: %v", err)
		}
		
		// 使用密钥解密事务数据
		txData, err := AESDecrypt(key, encTx.Ciphertext)
		if err != nil {
			return nil, fmt.Errorf("failed to decrypt transaction: %v", err)
		}
		
		decryptedTxs[i] = txData
	}
	
	return decryptedTxs, nil
}
